import logging

import docker
import yaml
from docker.errors import NotFound
from flask import Blueprint, g, request
from sqlalchemy import or_

from app.core.api import api_fail, api_success
from app.core.tools import model2dict
from app.extensions import cache, db
from app.models.admin import Config
from app.models.docker import DockerResource, DockerRunner
from app.tasks.vulnerability import start_vuln_resource, sync_remote_vulnerability_repo

logger = logging.getLogger(__name__)
bp = Blueprint("admin_vuln", __name__, url_prefix="/api/admin")


@bp.get("/vulnerability")
def vuln_list():
    """
    漏洞列表
    """
    page = int(request.args.get("page", 1))
    page_size = int(request.args.get("page_size", 10))
    search = request.args.get("search")
    db_query = db.session.query(DockerResource).filter(DockerResource.resource_type == "VUL")
    if search:
        db_query = db_query.filter(
            or_(
                DockerResource.name.contains(search)
                | DockerResource.app.contains(search)
                | DockerResource.cve.contains(search)
            )
        )
    page = db_query.order_by(DockerResource.updated_at.desc()).paginate(
        page=page, per_page=page_size
    )
    data = []
    for item in page.items:
        info = model2dict(item)
        info["docker_type_name"] = item.docker_type_name
        info["status_name"] = item.status_name
        # 判断key 是否存在
        if cache.exists("DOCKER_RESOURCE_%s" % item.id):
            info["building"] = True
        else:
            info["building"] = False
        data.append(info)
    return api_success({"total": page.total, "data": data})


@bp.get("/vulnerability/<int:pk>")
def vuln_detail(pk):
    """
    漏洞列表
    """
    item = DockerResource.get_by_id(pk)
    return api_success({"data": model2dict(item)})


@bp.delete("/vulnerability/<int:pk>")
def vuln_delete(pk):
    instance = DockerResource.get_by_id(pk)
    instance.delete()
    return api_success()


@bp.put("/vulnerability/<int:pk>")
def vuln_update(pk):
    """
    漏洞列表
    """
    item = DockerResource.get_by_id(pk)
    data = request.get_json()
    item.description = data["description"]
    item.docker_type = data["docker_type"]
    item.image = data["image"]
    # 检查镜像是否存在
    client = docker.DockerClient(Config.get_config(Config.KEY_DOCKER_API))
    if data.get("image"):
        try:
            client.images.get(data["image"])
            item.status = DockerResource.STATUS_BUILD
        except NotFound:
            item.status = DockerResource.STATUS_INIT
    else:
        item.status = DockerResource.STATUS_INIT
    item.name = data["name"]
    item.cve = data.get("cve", [])
    item.app = data.get("app")
    item.save()

    return api_success({"data": model2dict(item)})


@bp.post("/vulnerability")
def vuln_create():
    """
    添加漏洞
    """
    data = request.get_json()
    description = data["description"]
    docker_type = data["docker_type"]
    image = data["image"]
    name = data["name"]
    cve = data.get("cve", [])
    app = data.get("app")
    DockerResource.create(
        description=description,
        docker_type=docker_type,
        image=image,
        name=name,
        resource_type="VUL",
        app=app,
        cve=cve,
    )
    return api_success({})


@bp.post("/vulnerability/<int:pk>/run")
def vuln_run(pk):
    """
    添加漏洞
    """
    try:
        start_vuln_resource(pk, user_id=None, admin_id=g.user.id)
    except ValueError as e:
        return api_fail(msg=str(e))
    return api_success({})


@bp.get("/vulnerability/runner")
def vuln_runner():
    """
    添加漏洞
    """
    page = int(request.args.get("page", 1))
    page_size = int(request.args.get("page_size", 10))
    db_query = db.session.query(DockerRunner).filter()
    page = db_query.paginate(page=page, per_page=page_size)
    data = []
    for item in page.items:
        info = model2dict(item)
        info["resource"] = model2dict(item.resource)
        if item.user_id:
            info["username"] = item.user.username
        else:
            info["username"] = item.admin.username
        info["container_id"] = item.container_id[:8]
        data.append(info)
    return api_success({"total": page.total, "data": data})


@bp.post("/vulnerability/runner/destroy")
def vulnerability_runner_destroy():
    rid = request.get_json().get("id")
    instance: DockerRunner = db.session.query(DockerRunner).get(rid)
    docker_api = Config.get_config(Config.KEY_DOCKER_API)
    if not instance:
        return api_fail(msg="资源不存在")
    client = docker.DockerClient(docker_api)
    try:
        docker_container = client.containers.get(instance.container_id)
        docker_container.stop()
        docker_container.remove()
    except NotFound:
        pass
    instance.delete()
    return api_success()


@bp.delete("/vulnerability/<int:pk>")
def resource_delete(pk):
    """
    销毁容器
    """
    instance: DockerRunner = db.session.query(DockerResource).get(pk)
    docker_api = Config.get_config(Config.KEY_DOCKER_API)

    client = docker.DockerClient(docker_api)
    try:
        image = client.images.get(instance.image)
        image.remove(force=True)
    except NotFound as e:
        logger.error(e)
    db.session.delete(instance)
    db.session.commit()
    return api_success({})


@bp.post("/vulnerability/import")
def vuln_import():
    file = request.files["file"]
    filename = file.filename
    ext = filename.split(".")[-1]
    if ext != "yaml":
        return api_fail(msg="请上传yaml文件格式")
    try:
        yaml_data = yaml.safe_load(file)
    except yaml.YAMLError:
        return api_fail(msg="Error while parsing YAML file")
    currentImage = [i[0] for i in db.session.query(DockerResource.image)]
    bulk_create = []
    # 获取当前主机镜像
    client = docker.DockerClient(Config.get_config(Config.KEY_DOCKER_API))
    docker_images = client.images.list()
    tags = []
    for im in docker_images:
        tags += im.tags
    for item in yaml_data:
        image = item.get("image")
        if not image:
            continue
        if image in currentImage:
            logger.info("Pass Image:{}".format(item["image"]))
            continue
        if ":" in image:
            image_diff = image
        else:
            image_diff = f"{image}:latest"
        if image_diff in tags:
            status = DockerResource.STATUS_BUILD
        else:
            status = DockerResource.STATUS_INIT
        bulk_create.append(
            DockerResource(
                resource_type="CTF",
                name=item["name"],
                app=item.get("app"),
                image=item["image"],
                cve=item.get("cve", []),
                status=status,
                description=item.get("description"),
            )
        )
        db.session.bulk_save_objects(bulk_create)
    return api_success()


@bp.post("vulnerability/sync_vulnerability")
def sync_vulnerability():
    """
    同步远程漏洞
    """
    remote_repo = Config.get_config(Config.KEY_REMOTE_VULNERABILITY_REPOSITORY)
    if not remote_repo:
        return api_fail(msg="未配置远程漏洞仓库")
    sync_remote_vulnerability_repo.apply_async(args=(remote_repo,), kwargs={"admin_id": g.user.id})
    return api_success(msg="任务已提交")
